#![allow(dead_code)]

use core::{
    cell::{Cell, UnsafeCell},
    future::Future,
    pin::Pin,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};

use pin_utils::pin_mut;
use alloc::vec::Vec;
use crate::alloc::boxed::Box;
use crate::{rt_thread_self, rt_thread_resume, rt_thread_suspend, rt_schedule};

pub struct Executor {
    in_block_on: Cell<bool>,
    tasks: UnsafeCell<Vec<Task>>,
}

static VTABLE: RawWakerVTable = {
    unsafe fn clone(p: *const ()) -> RawWaker {
        RawWaker::new(p, &VTABLE)
    }
    unsafe fn wake(p: *const ()) {
        wake_by_ref(p)
    }
    unsafe fn wake_by_ref(p: *const ()) {
        (*(p as *const RawContext)).state.store(true, Ordering::Release);
        rt_thread_resume((*(p as *const RawContext)).thread);
    }
    unsafe fn drop(_: *const ()) {
        // no-op
    }

    RawWakerVTable::new(clone, wake, wake_by_ref, drop)
};

pub struct RawContext {
    thread: *mut usize,
    state: AtomicBool
}

pub(crate) static RESUME: AtomicBool = AtomicBool::new(false);

impl Executor {
    pub fn new() -> Self {
        Self {
            in_block_on: Cell::new(false),
            tasks: UnsafeCell::new(Vec::new()),
        }
    }

    pub fn block_on<T>(&self, f: impl Future<Output = T>) -> T {
        if self.in_block_on.get() {
            panic!()
        }
        self.in_block_on.set(true);

        pin_mut!(f);
        let ready = unsafe{ RawContext { thread: rt_thread_self(), state: AtomicBool::new(true) }};

        let waker =
            unsafe { Waker::from_raw(RawWaker::new(
                &ready as *const _ as *const _, &VTABLE)) };
        let val = loop {
            let mut task_woken = false;

            if ready.state.load(Ordering::Acquire) {
                task_woken = true;
                ready.state.store(false, Ordering::Release);

                let mut cx = Context::from_waker(&waker);
                if let Poll::Ready(val) = f.as_mut().poll(&mut cx) {
                    break val;
                }
            }
            let len = unsafe { (*self.tasks.get()).len() }; // (A)
            for i in 0..len {
                let task = unsafe { (*self.tasks.get()).get_unchecked(i) };

                if task.ready.state.load(Ordering::Acquire) {
                    task_woken = true;
                    task.ready.state.store(false, Ordering::Release);
                    let waker = unsafe {
                        Waker::from_raw(RawWaker::new(&task.ready as *const _ as *const _, &VTABLE))
                    };
                    let mut cx = Context::from_waker(&waker);
                    if unsafe {
                        !Pin::new_unchecked(&mut *task.raw.func.get())
                            .poll(&mut cx)
                            .is_ready()
                    } {
                        continue;
                    }
                }
            }
            if task_woken {
                continue;
            }
            if RESUME.load(Ordering::Acquire) {
                RESUME.store(false, Ordering::Release);
                continue;
            }

            unsafe {
                rt_thread_suspend(rt_thread_self());
                rt_schedule();
            };
        };
        self.in_block_on.set(false);
        val
    }

    pub fn spawn(&self, f: impl Future<Output = !> + 'static) {
        unsafe { (*self.tasks.get()).push(Task::new(f)) };
    }
}

pub struct Task {
    ready: RawContext,
    raw: RawTask<dyn Future<Output = !> + 'static>
}

struct RawTask<T> where T: ?Sized {
    func: Pin<Box<UnsafeCell<T>>>
}

impl Task {
    fn new(f: impl Future<Output = !> + 'static) -> Task {
        Task {
            ready: RawContext {
                thread: unsafe{rt_thread_self()},
                state: AtomicBool::new(true),
            },
            raw: RawTask{ func: Box::pin(UnsafeCell::new(f)) }
        }
    }
}